EMNLP 2023 | SAMP：基于自适应混合精度的训练后量化模型推理库
本文内容源自发表于计算机人工智能领域顶会 EMNLP 2023（The 2023 Conference on Empirical Methods in Natural Language Processing），由快手、腾讯等单位联合共同完成。
SAMP: A Model Inference Toolkit of Post-Training Quantization for Text Processing via Self-Adaptive Mixed-Precision
论文作者：
论文链接：
三分钟读论文
Powered by AI © PaperWeekly
摘要
AI 进入大模型时代，模型的计算性能是一个长期被关注的问题。目前工业界推理引擎已经验证半精度浮点（FP16）和 8 位整型（INT8）计算可以大大缩减计算位宽，从而加快模型推理速度。但已知的 INT8 量化方法使用门槛较高，且精度损失大，会导致模型效果大打折扣。我们推出了一个自适应混合精度推理库（SAMP），通过混合精度架构自动控制浮点型运算和整型运算，使得模型在量化推理加速的同时，保证计算精度。
文本理解是自然语言处理（NLP）领域的基本任务之一，例如在信息检索、对话系统、情感识别、摘要生成、语言模型等多个应用场景中。基于 Transformer 的模型在一系列下游任务中不断刷新着最佳记录（例如 BERT、XLNet、Google T5 等）。然而在一些大型工业系统中，由于在推理阶段缺乏高 GPU 占用考虑和良好的内存管理，训练框架（例如 TensorFlow 或 PyTorch）并不是模型部署的最佳选择。
用于深度学习模型推理加速的传统工具（例如 NVIDIA TensorRT、TurboTransformers 和 LightSeq）主要针对固定大小或可变长度输入而设计。这些工具的优化概念主要考虑在线计算系统中的内存管理、算子融合和一些数据剪枝技术。
计算精度方面，大多采用单一精度计算（使用纯浮点数计算或纯整型计算），加速性能有限。其中比较亮眼的是，NVIDIA 开发的 FasterTransformer 对 Transformer 模型进行了定点加速（将所有 Transformer 层使用 INT8 量化），与浮点计算相比取得了优异的加速效果。但这种在所有 Transformer 层进行全 INT8 量化的方式，模型推理结果很难达到原有浮点计算的精度，在很多任务中精度损失大，难广泛应用。
为了解决这些问题，我们提出了一种推理工具包 SAMP，其中包含自适应混合精度编码器和一系列先进的融合策略。自适应混合精度编码器可以在大量的通用矩阵乘法（GEMM）运算和 Transformer 层中找到最佳浮点定点混合精度组合方式，使模型推理的性能最贴近用户需求（计算精度或推理效率）。最终，混合精度计算取得了比全定点计算更好的计算精度。融合策略对 embedding 算子和量化相关计算操作进行融合改进，使得 CUDA 内核调用减少一半。
与此同时，SAMP 是由 C++ 编程语言实现的端到端工具包（从 Tokenizer 到 Embedding、Encoder 及下游任务），具有出色的推理速度，也降低了训练后量化推理的工业应用门槛。表 1 显示了与同类系统相比的创新点。
以下是 SAMP 工作的主要贡献（核心亮点）：
1. 自适应：SAMP 在训练后量化推理方法中平衡计算精度和延迟性能。用户可以针对不同的任务选择合适精度和推理延迟的混合精度配置。SAMP 还可通过自适应分配方法推荐给用户最佳的量化组合模式。
2. 推理效率：在较宽的精度范围（浮点到定点）中，SAMP 显示出比其他推理工具包更好的推理加速。在中文语言理解测评基准（CLUE）分类任务数据集中，与 FasterTransformer 相比，SAMP 实现了高达 1.05-1.15 倍的加速。
3. 灵活性：SAMP 涵盖众多下游任务，如分类、序列标记、文本匹配等。Target 模块是可扩展的并且可以灵活定制。它对用户友好且对平台依赖性较低。SAMP 支持 C++ 和 Python API，仅需要 CUDA 11.0 或更高版本即可。另外，SAMP 也提供了许多模型转换工具，支持不同格式模型之间相互转换。
SAMP 端到端推理库由四大模块组成：分词器（Tokenizer）、词嵌入层（Embedding）、编码器（Encoder）、下游目标层（Downstream Target）。各模块有不同亮点：
Tokenizer：支持中英文的多粒度分词（如：字符级、词语级以及常见的BertTokenizer），该模块通过 C++ 多线程并行加速，计算速度很快；
Embedding：SAMP 将 token_embedding、position_embedding、segment_embedding 三个 kernel 融合为一个 kernel，减少了 CUDA 的调用次数。
Encoder：在 Transformer 编码器里，SAMP 会通过精度衰减感知算法推荐出一个能使模型推理计算效率和精度兼得的混合精度配置。
Downstream Target：SAMP 支持各种各样的 NLP 下游任务。
SAMP 提供两种混合精度推理模式：全量化模式（Fully-Quant mode）和FFN量化模式（Quant-FFN-Only）。
1. 全量化模式（Fully-Quant mode）
全量化模式下，计算 Transformer 的 CUDA kernel 内部通过量化、反量化以及算子融合等方法，将所有 kernel 间的数据流动转为 8 位整型，大大缩减了传输位宽，推理速度达到了最快。
2. FFN量化模式（Quant-FFN-Only）
Quant-FFN-Only 模式是仅对 Transformer 中的 Feed-Forward 子网络做 8 位整型量化，保留了 Multi-Head-Attention 网络的浮点数计算（FP32 或 FP16）。这种方式因为精度损失较为平滑，在精度与推理效率的置换过程中，更容易被 SAMP 推荐。
3. SAMP混合精度推理库的使用方法
当用户明确给出模型精度要求或时延要求时：SAMP会给出最接近用户需求的混合精度推理参数和对应模式，进行计算。
用户未明确给出模型精度要求或时延要求时：SAMP将通过下图算法1中的精度衰减感知算法，自动推荐一个混合精度模式，进行计算。
成果展示
1. 绝对速度
我们测试了 SAMP 推理库在“全精度浮点（FP32）”、“半精度浮点（FP16）”、“全定点（INT8）”三种精度下的模型计算效率，选取主流的 PyTorch 和最新版 FasterTransformer 进行比较。
我们选取了较为常见的 batch_size 和文本长度组合。实验结果来看，大多情况下 SAMP 的计算效率是比 Nvidia 官方推理库 FasterTransformer 更快。
2. 精度、加速效果兼得
另外，我们通过中文语言理解测评基准（CLUE），选取一些文本分类任务——AFQMC（蚂蚁金服问答库）、IFLY-TEK（长文本分类）和 TNEWS（新闻短文本分类）等，对 SAMP 做了“计算精度 vs 计算加速比”的置换实验。
相比 PyTorch-FP16 和 FasterTransformer-FP16/INT8，SAMP 提供了 SAMP-FP16、SAMP-Fully-Quant 和 SAMP-Quant-FFN-Only 三种模式，其中后两种模式实现了 Transformer 层维度的混合精度计算。
上表中下划线的结果表示精度衰减感知算法推荐给用户的混合精度结果。以 AFQMC 任务为例，相比 PyTorch-FP16，SAMP-Fully-Quant 模式的模型精度从 0.7337 降至 0.6874，推理速度提升达到 4.7751 倍，FasterTransformer-INT8 的模型精度从 0.7337 降至 0.5773，推理速度提升只达到 3.4990 倍。同样，在另外两个实验结果不难得出类似结论—— SAMP 在“计算精度 vs 计算加速比”置换实验中取得了不错的效果。
1. 解决了现有后量化（PTQ）推理工具在工业应用中精度损失大的问题；
2. 推动了后量化（PTQ）技术在 NLP 多个下游任务中大规模使用；
3. 推理库轻量、灵活，对用户友好，支持用户自定义任务目标。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧